from gerrychain import Graph, Partition
import matplotlib.pyplot as plt
import geopandas as gpd
import ast

graph = Graph.from_json("../jsons/northcarolina.json")
gdf = gpd.read_file("../jsons/NC_VTD/NC_VTD.shp")


assignment_file = open("../outputs/NCoutput/PG_example_assignments.txt", "r")
pg_scores = open("../outputs/NCoutput/PG_example_scores.txt", "r")
assignment_lines = assignment_file.readlines()
pg_lines = pg_scores.readlines()


assignments = []
for i, line in enumerate(assignment_lines):
    assignment = ast.literal_eval(line)
    assignments.append(assignment)
    pg = pg_lines[i]
    print(f"PG score for example {i} is {pg}")
    partition = Partition(graph, assignment=assignment)
    ax = partition.plot(gdf)
    ax.axis("off")
    plt.savefig(f"../plots/High_Res_Plots/NC_{i}.png", dpi=1200)